"""
python fujian2_bp_predict.py
"""

import os
import json
import pandas as pd
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import re

# 设置输出路径
output_json_folder = './test/json_output'
output_graph_folder = './test/gragh_output'

# 创建输出文件夹
os.makedirs(output_json_folder, exist_ok=True)
os.makedirs(output_graph_folder, exist_ok=True)

# 读取 JSON 数据
def read_json_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df['sales'] = df['sales'].astype(int)
    return df

# 准备训练数据
def prepare_data(df):
    df = df.sort_values('date')
    df.set_index('date', inplace=True)
    df = df.asfreq('D', fill_value=0)  # 填充缺失的日期
    X = np.array(df.index.astype(np.int64) // 10**9).reshape(-1, 1)  # 转换为时间戳
    y = df['sales'].values
    return X, y

# 构建 BP 神经网络模型
def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(1,)),
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dense(1)  # 输出层
    ])
    model.compile(optimizer='adam', loss='mse')
    return model

# 进行预测
def predict_future_sales(model, last_date, num_days):
    future_dates = [last_date + timedelta(days=i) for i in range(1, num_days + 1)]
    future_dates_timestamp = np.array([int(date.timestamp()) for date in future_dates]).reshape(-1, 1)  # 转换为时间戳
    predictions = model.predict(future_dates_timestamp)
    return future_dates, predictions.flatten()


# 保存预测结果到 JSON 文件
def save_predictions_to_json(category_id, future_dates, predicted_sales):
    result = [{'date': date.strftime('%Y/%m/%d'), 'predicted_sales': int(sale)} for date, sale in zip(future_dates, predicted_sales)]
    output_json_path = os.path.join(output_json_folder, f'json_category{category_id}.json')
    with open(output_json_path, 'w') as f:
        json.dump(result, f)

# 绘制原始数据和预测数据
def plot_data_and_predictions(original_dates, original_sales, future_dates, predicted_sales, category_id):
    plt.figure(figsize=(10, 6))
    plt.scatter(original_dates, original_sales, label='Original Data', color='blue')
    plt.scatter(future_dates, predicted_sales, label='Predicted Data', color='red')
    plt.legend()
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title(f'BP Neural Network Prediction on Category {category_id}')
    output_image_path = os.path.join(output_graph_folder, f'gragh_category{category_id}.png')
    plt.savefig(output_image_path)
    plt.close()

# 主程序
if __name__ == "__main__":
    input_folder = r'..\fujian\fujian2\groupByCategory'
    json_files = [os.path.join(input_folder, file) for file in os.listdir(input_folder) if file.endswith('.json')]

    for json_file in json_files:
        filename = os.path.basename(json_file)
        print(filename)
        # match = re.search(r'category_(\d+)', filename)
        match = re.search(r'category_category(\d+)', filename)
        if match:
            category_id = int(match.group(1))
        else:
            print("未找到匹配的 category ID")
            continue

        # 读取数据
        df = read_json_data(json_file)
        
        # 准备数据
        X, y = prepare_data(df)

        # 构建和训练模型
        model = build_model()
        model.fit(X, y, epochs=100, verbose=0)

        # 预测未来销售
        last_date = df['date'].max()
        future_dates, predicted_sales = predict_future_sales(model, last_date, 92)

        # 保存预测到 JSON 文件
        save_predictions_to_json(category_id, future_dates, predicted_sales)

        # 绘制原始数据和预测数据的散点图
        plot_data_and_predictions(df['date'], df['sales'], future_dates, predicted_sales, category_id)
